import os
import numpy as np
import math


def get_pic_result(pic_list):
    pass


def get_file(pic_list, class_names):
    """
    Args:
        pic_list: include n_classes pics path
    Returns:
        list of images and labels
    """
    n_classes = len(pic_list)
    classes = []
    labels = []
    name_to_label = {}
    for i in range(n_classes):
        name_to_label[class_names[i]] = i
        classes.append(pic_list[i])
        label = []
        for j in range(len(pic_list[i])):
            label.append(i)
        labels.append(label)

    for i in range(n_classes):
        print(f'there are {len(labels[i])} {class_names[i]}s')

    image_list = []
    label_list = []
    for i in range(n_classes):
        image_list = np.hstack((image_list, classes[i]))
        label_list = np.hstack((label_list, labels[i]))

    temp = np.array([image_list, label_list])
    temp = temp.transpose()
    np.random.shuffle(temp)

    all_image_list = temp[:, 0]
    all_label_list = temp[:, 1]

    n_sample = len(all_label_list)
    n_val = math.ceil(n_sample * 0.2)  # number of validation samples
    n_train = n_sample - n_val  # number of trainning samples

    tra_images = all_image_list[0:n_train]
    tra_labels = all_label_list[0:n_train]
    tra_labels = [int(float(i)) for i in tra_labels]
    val_images = all_image_list[n_train:-1]
    val_labels = all_label_list[n_train:-1]
    val_labels = [int(float(i)) for i in val_labels]
    print(tra_images)
    print(tra_labels)
    print(val_images)
    print(val_labels)

    return tra_images, tra_labels, val_images, val_labels


def img_to_list(file_dir):
    cats = []
    dogs = []
    for file in os.listdir(file_dir):
        name = file.split(sep='.')
        if name[0] == 'cat':
            cats.append(file_dir + file)
        elif name[0] == 'dog':
            dogs.append(file_dir + file)
    pic_list = [cats, dogs]
    return pic_list


if __name__ == '__main__':
    file_dir = 'img/'
    img_list = img_to_list(file_dir)
    # get_file(img_list, ['cat', 'dog'])
